import tensorflow as tf
import numpy as np
import os
from PIL import Image
import random


class CNN(object):
    def __init__(self, image_height, image_width, max_captcha, char_set, model_save_dir):
        # 初始值
        self.image_height = image_height
        self.image_width = image_width
        self.max_captcha = max_captcha
        self.char_set = char_set
        self.char_set_len = len(char_set)
        self.model_save_dir = model_save_dir  # 模型路径
        with tf.compat.v1.name_scope('parameters'):
            self.w_alpha = 0.01
            self.b_alpha = 0.1
        # tf.compat.v1初始化占位符
        with tf.compat.v1.name_scope('data'):
            self.X = tf.compat.v1.placeholder(tf.compat.v1.float32, [None, self.image_height * self.image_width])  # 特征向量
            self.Y = tf.compat.v1.placeholder(tf.compat.v1.float32, [None, self.max_captcha * self.char_set_len])  # 标签
            self.keep_prob = tf.compat.v1.placeholder(tf.compat.v1.float32)  # dropout值

    @staticmethod
    def convert2gray(img):
        """
        图片转为灰度图，如果是3通道图则计算，单通道图则直接返回
        :param img:
        :return:
        """
        if len(img.shape) > 2:
            r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
            gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
            return gray
        else:
            return img

    def text2vec(self, text):
        """
        转标签为oneHot编码
        :param text: str
        :return: numpy.array
        """
        text_len = len(text)
        if text_len > self.max_captcha:
            raise ValueError('验证码最长{}个字符'.format(self.max_captcha))

        vector = np.zeros(self.max_captcha * self.char_set_len)
        # 验证码数量*验证码字符集的行向量
        for i, ch in enumerate(text):
            idx = i * self.char_set_len + self.char_set.index(ch)
            vector[idx] = 1
            # 验证码对应的值为1，如验证码1234，字符集为{1，2，3，4}
            # 则对应的向量为[1,0,0,0,|0,1,0,0,|0,1,0,0,|0,0,0,1]
            # |为分隔符
        return vector

    def model(self):
        x = tf.compat.v1.reshape(self.X, shape=[-1, self.image_height, self.image_width, 1])
        # 将tensor变换为参数shape的形式。第一个参数表示值张量Tensor，这里self.X为特征向量
        # 第二这个参数也是一个张量Tensor，用于定义输出张量的形状。
        # -1表示第一维未知，self.image_height, self.image_width分别代表图片长和宽，1表示黑白图片
        print(">>> input x: {}".format(x))

        # 卷积层1
        wc1 = tf.compat.v1.get_variable(name='wc1', shape=[3, 3, 1, 32], dtype=tf.compat.v1.float32,
                              initializer=tf.compat.v1.keras.initializers.glorot_normal())
        # 创建tensorflow变量，名为wc1，shape（变量维度）为3*3*1*32，
        # 3*3为卷积核的大小，1为矩阵深度，32为卷积核的数量
        # dtype（变量类型）为浮点数
        # initializer（变量初始化方式）
        # tf.compat.v1.keras.initializers.glorot_normal()是一种带权重的初始化方法，旨在使所有图层的渐变比例大致相同
        bc1 = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([32]))
        # b_alpha是参数中的一个变量，默认值为0.1，
        # tf.compat.v1.random.normal从服从指定正态分布的序列中随机取出32个数
        # 得到的bc1为行向量
        conv1 = tf.compat.v1.nn.relu(tf.compat.v1.nn.bias_add(tf.compat.v1.nn.conv2d(x, wc1, strides=[1, 1, 1, 1], padding='SAME'), bc1))
        # tf.compat.v1.nn.relu为激励函数，得到卷积后的结果
        # tf.compat.v1.nn.conv2d为二维卷积，x为张量输入，wc1为卷积核
        # strides为步长，第一项和第四项必须为1，第二三项分别别为水平步长和垂直步长
        # padding为填充，padding='SAME'代表全零填充
        # tf.compat.v1.nn.bias_add代表在卷积后的图像的每一行加上bc1，其中bc1为偏差项
        conv1 = tf.compat.v1.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        # 池化操作，conv1为上一步计算结果
        # ksize：池化窗口的大小，是一个长度为4的一维数组，
        # 但是这个数组的第一个和最后一数必须为1，即[1, height, width, 1]。
        # 这意味着池化层的过滤器是不可以在batch和channels上做池化。
        # 实际应用中，使用最多的过滤器尺寸为[1, 2, 2, 1]或者[1, 3, 3, 1]。
        # strides：窗口每一个维度上滑动步长，
        # padding：填充
        conv1 = tf.compat.v1.nn.dropout(conv1, self.keep_prob)
        # tf.compat.v1.nn.dropout()是tensorflow里面为了防止或减轻过拟合而使用的函数，self.keep_prob代表元素被丢弃的概率
        # 卷积层2
        wc2 = tf.compat.v1.get_variable(name='wc2', shape=[3, 3, 32, 64], dtype=tf.compat.v1.float32,
                              initializer=tf.compat.v1.keras.initializers.glorot_normal())
        bc2 = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([64]))
        conv2 = tf.compat.v1.nn.relu(tf.compat.v1.nn.bias_add(tf.compat.v1.nn.conv2d(conv1, wc2, strides=[1, 1, 1, 1], padding='SAME'), bc2))
        conv2 = tf.compat.v1.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv2 = tf.compat.v1.nn.dropout(conv2, self.keep_prob)

        # 卷积层3
        wc3 = tf.compat.v1.get_variable(name='wc3', shape=[3, 3, 64, 128], dtype=tf.compat.v1.float32,
                              initializer=tf.compat.v1.keras.initializers.glorot_normal())
        bc3 = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([128]))
        conv3 = tf.compat.v1.nn.relu(tf.compat.v1.nn.bias_add(tf.compat.v1.nn.conv2d(conv2, wc3, strides=[1, 1, 1, 1], padding='SAME'), bc3))
        conv3 = tf.compat.v1.nn.max_pool(conv3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv3 = tf.compat.v1.nn.dropout(conv3, self.keep_prob)
        print(">>> convolution 3: ", conv3.shape)
        # next_shape = conv3.shape[1] * conv3.shape[2] * conv3.shape[3]
        # 卷积层4
        wc4 = tf.compat.v1.get_variable(name='wc4', shape=[3, 3, 128, 128], dtype=tf.compat.v1.float32,
                              initializer=tf.compat.v1.keras.initializers.glorot_normal())
        bc4 = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([128]))
        conv4 = tf.compat.v1.nn.relu(tf.compat.v1.nn.bias_add(tf.compat.v1.nn.conv2d(conv3, wc4, strides=[1, 1, 1, 1], padding='SAME'), bc4))
        # conv3 = tf.compat.v1.nn.max_pool(conv3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv4 = tf.compat.v1.nn.dropout(conv4, self.keep_prob)
        print(">>> convolution 4: ", conv4.shape)
        # next_shape = conv4.shape[1] * conv4.shape[2] * conv4.shape[3]
        wc5 = tf.compat.v1.get_variable(name='wc5', shape=[3, 3, 128, 128], dtype=tf.compat.v1.float32,
                                        initializer=tf.compat.v1.keras.initializers.glorot_normal())
        bc5 = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([128]))
        conv5 = tf.compat.v1.nn.relu(
            tf.compat.v1.nn.bias_add(tf.compat.v1.nn.conv2d(conv4, wc5, strides=[1, 1, 1, 1], padding='SAME'), bc5))
        # conv3 = tf.compat.v1.nn.max_pool(conv3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv5 = tf.compat.v1.nn.dropout(conv5, self.keep_prob)
        print(">>> convolution 5: ", conv5.shape)
        next_shape = conv5.shape[1] * conv5.shape[2] * conv5.shape[3]

        # 全连接层1
        wd1 = tf.compat.v1.get_variable(name='wd1', shape=[next_shape, 1024], dtype=tf.compat.v1.float32,
                              initializer=tf.compat.v1.keras.initializers.glorot_normal())
        bd1 = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([1024]))
        dense = tf.compat.v1.reshape(conv5, [-1, wd1.get_shape().as_list()[0]])
        dense = tf.compat.v1.nn.relu(tf.compat.v1.add(tf.compat.v1.matmul(dense, wd1), bd1))
        dense = tf.compat.v1.nn.dropout(dense, self.keep_prob)

        # 全连接层2
        wout = tf.compat.v1.get_variable('name', shape=[1024, self.max_captcha * self.char_set_len], dtype=tf.compat.v1.float32,
                               initializer=tf.compat.v1.keras.initializers.glorot_normal())
        bout = tf.compat.v1.Variable(self.b_alpha * tf.compat.v1.random.normal([self.max_captcha * self.char_set_len]))
        #
        with tf.compat.v1.name_scope('y_prediction'):
            y_predict = tf.compat.v1.add(tf.compat.v1.matmul(dense, wout), bout)
            # matmul代表相乘,add表示相加
        return y_predict
